/**
* This file is part of LIO-mapping.
* 
* Copyright (C) 2019 Haoyang Ye <hy.ye at connect dot ust dot hk>,
* Robotics and Multiperception Lab (RAM-LAB <https://ram-lab.com>),
* The Hong Kong University of Science and Technology
* 
* For more information please see <https://ram-lab.com/file/hyye/lio-mapping>
* or <https://sites.google.com/view/lio-mapping>.
* If you use this code, please cite the respective publications as
* listed on the above websites.
* 
* LIO-mapping is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
* 
* LIO-mapping is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
* 
* You should have received a copy of the GNU General Public License
* along with LIO-mapping.  If not, see <http://www.gnu.org/licenses/>.
*/

//
// Created by hyye on 3/27/18.
//

#ifndef LIO_ESTIMATOR_H_
#define LIO_ESTIMATOR_H_

#include<Eigen/StdVector>

#include "imu_processor/MeasurementManager.h"
#include "imu_processor/IntegrationBase.h"
#include "imu_processor/ImuInitializer.h"
#include "point_processor/PointMapping.h"

#include "factor/PoseLocalParameterization.h"
#include "factor/GravityLocalParameterization.h"
#include "factor/ImuFactor.h"
#include "factor/PointDistanceFactor.h"
#include "factor/PlaneProjectionFactor.h"
#include "factor/PriorFactor.h"
#include "factor/MarginalizationFactor.h"
#include "factor/PlaneToPlaneFactor.h"
#include "factor/PivotPointPlaneFactor.h"

#include <std_srvs/SetBool.h>

#include "visualizer/Visualizer.h"
#include "imu_processor/estimator_trajectory_manager.h"
#include "visualizer/splineVisualizer.h"
#include <deque>

//#define FIX_MAP
//#define USE_CORNER

namespace lio {

using Eigen::Vector3d;
using std::shared_ptr;
using std::unique_ptr;

typedef multimap<float, pair<PointT, PointT>, greater<float> > ScorePointCoeffMap;

enum EstimatorStageFlag {
  NOT_INITED,
  INITED,
};

struct StampedTransform {
  double time;
  Transform transform;
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
};

struct EstimatorConfig {
  size_t window_size = 15;
  size_t opt_window_size = 5;
  int init_window_factor = 3;
  int estimate_extrinsic = 2;

  float corner_filter_size = 0.2;
  float surf_filter_size = 0.4;
  float map_filter_size = 0.6;

  float min_match_sq_dis = 1.0;
  float min_plane_dis = 0.2;
  Transform transform_lb{Eigen::Quaternionf(1, 0, 0, 0), Eigen::Vector3f(0, 0, -0.1)};

  bool opt_extrinsic = false;

  bool run_optimization = true;
  bool update_laser_imu = true;
  bool gravity_fix = true;
  bool plane_projection_factor = true;
  bool imu_factor = true;
  bool point_distance_factor = false;
  bool prior_factor = false;
  bool marginalization_factor = true;
  bool pcl_viewer = false;

  bool enable_deskew = true; ///< if disable, deskew from PointOdometry will be used
  bool cutoff_deskew = false;
  bool keep_features = false;

  double gry_weight = 1;
  double acc_weight = 1;
  double lidar_weight = 1;
  double ba_weight = 1;
  double bg_weight = 1;
  Eigen::Vector3d gry_bias = Eigen::Vector3d::Zero();
  Eigen::Vector3d acc_bias = Eigen::Vector3d::Zero();

  double knot_distance = 0.1;
  int pre_knot_size = 10;
  int surroundingKeyframeSearchNum = 120;
  int iteration_size = 1;
  bool use_first_imu = false;

  IntegrationBaseConfig pim_config;
};

class Estimator : public MeasurementManager, public PointMapping {
 public:
  Estimator();
  Estimator(EstimatorConfig config, MeasurementManagerConfig mm_config = MeasurementManagerConfig());
  ~Estimator();
  void ClearState();
  void SetupRos(ros::NodeHandle &nh);

  void SetupAllEstimatorConfig(const EstimatorConfig &config,
                               const MeasurementManagerConfig &mm_config);

  void ProcessEstimation();
  void ProcessImu(double dt,
                  const Vector3d &linear_acceleration,
                  const Vector3d &angular_velocity);
  void ProcessLaserOdom(const Transform &transform_in, const std_msgs::Header &header);
  void ProcessCompactData(const sensor_msgs::PointCloud2ConstPtr &compact_data, const std_msgs::Header &header);

  void BuildLocalMap(vector<FeaturePerFrame> &feature_frames);
  // for spline
  void BuildLocalMapSplineInit(vector<FeaturePerFrame> &feature_frames); // 初始化时,构建局部地图
  void BuildLocalMapSplineNormal(vector<FeaturePerFrame> &feature_frames,
                                Eigen::Vector3d &pos_imu_pivot, 
                                Eigen::Quaterniond &rot_imu_pivot, bool first_opt = true); // 初始化完成后,构建局部地图

#ifdef USE_CORNER
  void CalculateFeatures(const pcl::KdTreeFLANN<PointT>::Ptr &kdtree_surf_from_map,
                         const PointCloudPtr &local_surf_points_filtered_ptr,
                         const PointCloudPtr &surf_stack,
                         const pcl::KdTreeFLANN<PointT>::Ptr &kdtree_corner_from_map,
                         const PointCloudPtr &local_corner_points_filtered_ptr,
                         const PointCloudPtr &corner_stack,
                         const Transform &local_transform,
                         vector<unique_ptr<Feature>> &features);

  void CalculateLaserOdom(const pcl::KdTreeFLANN<PointT>::Ptr &kdtree_surf_from_map,
                          const PointCloudPtr &local_surf_points_filtered_ptr,
                          const PointCloudPtr &surf_stack,
                          const pcl::KdTreeFLANN<PointT>::Ptr &kdtree_corner_from_map,
                          const PointCloudPtr &local_corner_points_filtered_ptr,
                          const PointCloudPtr &corner_stack,
                          Transform &local_transform,
                          vector<unique_ptr<Feature>> &features);
#else
  void CalculateFeatures(const pcl::KdTreeFLANN<PointT>::Ptr &kdtree_surf_from_map,
                         const PointCloudPtr &local_surf_points_filtered_ptr,
                         const PointCloudPtr &surf_stack,
                         const Transform &local_transform,
                         vector<unique_ptr<Feature>> &features);

  void CalculateLaserOdom(const pcl::KdTreeFLANN<PointT>::Ptr &kdtree_surf_from_map,
                          const PointCloudPtr &local_surf_points_filtered_ptr,
                          const PointCloudPtr &surf_stack,
                          Transform &local_transform,
                          vector<unique_ptr<Feature>> &features);
#endif

  void SolveOptimization();

  void SlideWindow();

  void VectorToDouble();
  void DoubleToVector();

  bool RunInitialization();

  bool SplineInitialization();

  bool SendCloudAndPose();

  bool SetNewContralPoint();

  bool GetLastPVQ();  // 获取当前优化后，最后一帧点云时间处对应的 IMU 的 p v q，用于预测 IMU 位姿

  void DeskewOneCloud(const PointCloudPtr &pi, double te);  // 对一帧点云做 transform to end
  void DeskewOneCloud(const PointCloudPtr &pi, PointCloudPtr &po, double te);  // 对一帧点云做 transform to end
  PointCloudPtr transformPointCloud(PointCloudPtr &cloudIn, double time);  // 转换已经去畸变的点云到 I_0 系
  void PointAssociateToMap(const PointT &pi, PointT &po, double time); // 转换一个点到 I_0 系
  // 构建局部地图
  void extractSurroundingKeyFrames();

  // 构建残差
  void cornerOptimization();
  void surfOptimization();
  void imuOptimization();

  // 发布点云
  void publishDeskewCloud();
  void publishGlobalMap();
  void publishPredictCloud();
  void publishNewestCloud();
  void publishLocalMap();
  void pubDeskewThread();
  void pubGlobalMapThread();
  void pubOneFrameDeskewPoints();

  void pubNewestControlPoint();
  void pubOptimizedControlPoint();

  void saveAcceAndGrysMeasurements();


  PlaneNormalVisualizer normal_vis;

  int extrinsic_stage_ = 2;

  PointCloudPtr local_surf_points_ptr_, local_surf_points_filtered_ptr_;
  PointCloudPtr local_corner_points_ptr_, local_corner_points_filtered_ptr_;

 protected:
  EstimatorStageFlag stage_flag_ = NOT_INITED;
  EstimatorConfig estimator_config_;

  bool first_imu_ = false;
  double initial_time_ = -1;

  size_t cir_buf_count_ = 0;
  size_t laser_odom_recv_count_ = 0;

  Vector3d acc_last_, gyr_last_;
  Vector3d g_vec_ = Vector3d(0.0, 0.0, -9.79);
  bool init_local_map_ = false;

  CircularBuffer<vector<Vector3d> > linear_acceleration_buf_{estimator_config_.window_size + 1};
  CircularBuffer<vector<Vector3d> > angular_velocity_buf_{estimator_config_.window_size + 1};
  // 这两个时间戳直接使用 Header 中的时间，不需要设置相对时间
  CircularBuffer<vector<double> > imu_timestamp_buf_{estimator_config_.window_size + 1};

  // 全部的历史 de-skew 点云
  std::deque<PointCloudPtr> past_surfs_, past_corners_;
  std::deque<double> past_timestamps_;
  // skew 点云
  std::deque<PointCloudPtr> skew_surfs_, skew_corners_, skew_fulls_;
  std::deque<double> skew_timestamps_;
  // 缓冲区点云
  std::deque<PointCloudPtr> buffer_surfs_, buffer_corners_, buffer_fulls_;
  std::deque<double> buffer_timestamps_;

  // 局部地图点云
  std::deque<PointCloudPtr> recent_surfs_, recent_corners_;

  // 存储原始数据
  std::queue<ImuMsgConstPtr> origin_imu_buf_;
  std::queue<CompactDataConstPtr> origin_compact_data_buf_;
  std::vector<Vector3d> newKnotImu_gry, newKnotImu_acc;
  std::vector<double> newKnotImu_time;

//  Transform transform_lb_{Eigen::Quaternionf(1, 0, 0, 0), Eigen::Vector3f(-0.05, 0, 0.05)}; ///< Base to laser transform
  Transform transform_lb_{Eigen::Quaternionf(1, 0, 0, 0), Eigen::Vector3f(0, 0, -0.1)}; ///< Base to laser transform

  Eigen::Matrix3d R_WI_; ///< R_WI is the rotation from the inertial frame into Lidar's world frame
  Eigen::Quaterniond Q_WI_; ///< Q_WI is the rotation from the inertial frame into Lidar's world frame

  tf::StampedTransform wi_trans_, laser_local_trans_, laser_predict_trans_, laser_send_trans_;
  tf::TransformBroadcaster tf_broadcaster_est_;

  ros::Publisher pub_predict_odom_;
  nav_msgs::Odometry predict_odom_;

  ros::Publisher pub_discrete_imu_odom_;
  nav_msgs::Odometry discrete_imu_odom_;

  ros::Publisher pub_newest_odom_;
  nav_msgs::Odometry newest_odom_;

  ros::Publisher pub_past_odom_;
  nav_msgs::Odometry past_odom_;

  ros::Publisher pub_control_point_, pub_opt_cp_;
  nav_msgs::Odometry control_point_;

  ros::Publisher pub_plane_normal_;
  visualization_msgs::MarkerArray plane_normal_array_;

  ros::Publisher pub_local_surf_points_;
  ros::Publisher pub_local_corner_points_;
  ros::Publisher pub_local_full_points_;

  ros::Publisher pub_past_surf_points_;
  ros::Publisher pub_past_corner_points_;
  ros::Publisher pub_past_full_points_;

  ros::Publisher pub_map_surf_points_;
  ros::Publisher pub_map_corner_points_;
  ros::Publisher pub_predict_surf_points_;
  ros::Publisher pub_predict_corner_points_;
  ros::Publisher pub_predict_full_points_;
  ros::Publisher pub_predict_corrected_full_points_;

  ros::Publisher pub_newest_surf_points_;
  ros::Publisher pub_newest_corner_points_;
  ros::Publisher pub_newest_full_points_;

  ros::Publisher pub_extrinsic_;

  Visualizer vis_bef_opt{"vis_bef_opt", vector<double>{0.0, 0.0, 1.0}, vector<double>{1.0, 1.0, 1.0}};
  Visualizer vis_aft_opt{"vis_aft_opt"};

  Vector3d P_pivot_;
  Matrix3d R_pivot_;

  bool convergence_flag_ = false;

  CircularBuffer<StampedTransform> imu_stampedtransforms{100};

 private:
//  double para_qwi_[SIZE_QUAT];
  double g_norm_;
  bool gravity_fixed_ = false;

  Transform transform_tobe_mapped_bef_;
  Transform transform_es_;

  // for spline
  // CircularBuffer<vector<sensor_msgs::ImuConstPtr>> imu_raw_stack_{estimator_config_.window_size + 1};
  // 创建 TrajectoryManager 类的智能指针
  TrajectoryManager::Ptr traj_manager_;
  double time_init_ = 0.0;
  bool is_get_time_init_ = false;

  struct imuData {
    double time;  // 时间戳
    Eigen::Vector3d w; // 角速度 w
    Eigen::Vector3d a; // 加速度 a
    Eigen::Vector3d p; // 位置	p
    Eigen::Quaterniond q; // 旋转	q
    Eigen::Vector3d v; // 速度	v
    imuData() : time(-1.0), w(0., 0., 0.),
                a(0., 0., 0.), p(0., 0., 0.),
                q(1., 0., 0., 0.), v(0., 0., 0.) {}
  };
  imuData imu_data_last_;
  bool system_init_ = false;
  bool built_traj_ = false;
  int last_recent_id = -1; // 局部地图中最新一帧的点云序号，用来判断是否需要更新局部地图
  PointCloudPtr laserCloudCornerFromMap, laserCloudSurfFromMap;
  PointCloudPtr laserCloudCornerFromMapDS, laserCloudSurfFromMapDS;
  pcl::KdTreeFLANN<PointT>::Ptr kdtree_surf_from_map, kdtree_corner_from_map;

  Eigen::Vector3d ba_last = Eigen::Vector3d(0, 0, 0);
  Eigen::Vector3d bg_last = Eigen::Vector3d(0, 0, 0);
  Transform lidar_pose_cur;

  splineVisualizer::Ptr splineVisual;
  Eigen::Quaterniond first_imu_q;
};

}

#endif //LIO_ESTIMATOR_H_
